import math

import torch.nn as nn

from .shufflenet import ShuffleNet
from .bifpn import BiFPN
from .head import Head


class CenterNet(nn.Module):
    def __init__(self, base_channel=32, repeats=3, num_cls=1) -> None:
        super().__init__()
        channels = [int(base_channel * math.pow(2, n)) for n in range(5)]
        self.backbone = ShuffleNet(channels_list=channels)
        self.neck = BiFPN(in_channels=channels[1:], out_channels=channels[1], repeats=repeats)
        self.head = Head(in_channel=channels[1], num_cls=num_cls)

    def forward(self, x):
        x = self.backbone(x)
        x = self.neck(x)
        x = self.head(x)
        return x
